use pest::Parser;
use pest_derive::Parser;

#[derive(Parser)]
#[grammar = "lexer.pest"]
struct SysYLexer;

#[derive(Debug, PartialEq)]
pub enum TokenType {
    Const,
    Int,
    Void,
    If,
    Else,
    While,
    Break,
    Continue,
    Return,
    Plus,
    Minus,
    Mul,
    Div,
    Mod,
    Assign,
    Eq,
    Neq,
    Lt,
    Gt,
    Le,
    Ge,
    Not,
    And,
    Or,
    LParen,
    RParen,
    LBrace,
    RBrace,
    LBrackt,
    RBrackt,
    Comma,
    Semicolon,
    Ident(String),
    IntegerConst(usize),
}

impl TokenType {
    pub fn stringify(&self) -> String {
        match self {
            TokenType::Const => "CONST".to_string(),
            TokenType::Int => "INT".to_string(),
            TokenType::Void => "VOID".to_string(),
            TokenType::If => "IF".to_string(),
            TokenType::Else => "ELSE".to_string(),
            TokenType::While => "WHILE".to_string(),
            TokenType::Break => "BREAK".to_string(),
            TokenType::Continue => "CONTINUE".to_string(),
            TokenType::Return => "RETURN".to_string(),
            TokenType::Plus => "PLUS".to_string(),
            TokenType::Minus => "MINUS".to_string(),
            TokenType::Mul => "MUL".to_string(),
            TokenType::Div => "DIV".to_string(),
            TokenType::Mod => "MOD".to_string(),
            TokenType::Assign => "ASSIGN".to_string(),
            TokenType::Eq => "EQ".to_string(),
            TokenType::Neq => "NEQ".to_string(),
            TokenType::Lt => "LT".to_string(),
            TokenType::Gt => "GT".to_string(),
            TokenType::Le => "LE".to_string(),
            TokenType::Ge => "GE".to_string(),
            TokenType::Not => "NOT".to_string(),
            TokenType::And => "AND".to_string(),
            TokenType::Or => "OR".to_string(),
            TokenType::LParen => "L_PAREN".to_string(),
            TokenType::RParen => "R_PAREN".to_string(),
            TokenType::LBrace => "L_BRACE".to_string(),
            TokenType::RBrace => "R_BRACE".to_string(),
            TokenType::LBrackt => "L_BRACKT".to_string(),
            TokenType::RBrackt => "R_BRACKT".to_string(),
            TokenType::Comma => "COMMA".to_string(),
            TokenType::Semicolon => "SEMICOLON".to_string(),
            TokenType::Ident(name) => format!("IDENT {}", name),
            TokenType::IntegerConst(value) => format!("INTEGER_CONST {}", value),
        }
    }
}

impl std::fmt::Display for TokenType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.stringify())
    }
}

#[derive(Debug)]
pub struct Token {
    token_type: TokenType,
    lexeme: String,
    line: usize,
}

impl std::fmt::Display for Token {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.token_type {
            TokenType::Ident(_) | TokenType::IntegerConst(_) => {
                write!(f, "{} at Line {}.", self.token_type, self.line)
            }
            _ => write!(
                f,
                "{} {} at Line {}.",
                self.token_type, self.lexeme, self.line
            ),
        }
    }
}

#[derive(Debug)]
struct ErrorMessage {
    line: usize,
    message: String,
}

#[derive(Debug)]
pub struct LexerError {
    messages: Vec<ErrorMessage>,
}

impl std::fmt::Display for LexerError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for msg in &self.messages {
            writeln!(f, "Error type A at Line {}: {}.", msg.line, msg.message)?;
        }
        Ok(())
    }
}

pub fn tokenize(input: &str) -> Result<Vec<Token>, LexerError> {
    let mut tokens: Vec<Token> = vec![];
    let mut errors: Vec<ErrorMessage> = vec![];
    let pairs = SysYLexer::parse(Rule::PROGRAM, input).unwrap_or_else(|e| panic!("{}", e));
    for pair in pairs.into_iter() {
        match pair.as_rule() {
            Rule::TOKEN => {
                for pair in pair.into_inner() {
                    let span = pair.as_span();
                    let lexeme = span.as_str().to_string();
                    let line = span.start_pos().line_col().0;
                    let token_type = match pair.as_rule() {
                        Rule::CONST => TokenType::Const,
                        Rule::INT => TokenType::Int,
                        Rule::VOID => TokenType::Void,
                        Rule::IF => TokenType::If,
                        Rule::ELSE => TokenType::Else,
                        Rule::WHILE => TokenType::While,
                        Rule::BREAK => TokenType::Break,
                        Rule::CONTINUE => TokenType::Continue,
                        Rule::RETURN => TokenType::Return,
                        Rule::PLUS => TokenType::Plus,
                        Rule::MINUS => TokenType::Minus,
                        Rule::MUL => TokenType::Mul,
                        Rule::DIV => TokenType::Div,
                        Rule::MOD => TokenType::Mod,
                        Rule::ASSIGN => TokenType::Assign,
                        Rule::EQ => TokenType::Eq,
                        Rule::NEQ => TokenType::Neq,
                        Rule::LT => TokenType::Lt,
                        Rule::GT => TokenType::Gt,
                        Rule::LE => TokenType::Le,
                        Rule::GE => TokenType::Ge,
                        Rule::NOT => TokenType::Not,
                        Rule::AND => TokenType::And,
                        Rule::OR => TokenType::Or,
                        Rule::LPAREN => TokenType::LParen,
                        Rule::RPAREN => TokenType::RParen,
                        Rule::LBRACE => TokenType::LBrace,
                        Rule::RBRACE => TokenType::RBrace,
                        Rule::LBRACKT => TokenType::LBrackt,
                        Rule::RBRACKT => TokenType::RBrackt,
                        Rule::COMMA => TokenType::Comma,
                        Rule::SEMICOLON => TokenType::Semicolon,
                        Rule::IDENT => TokenType::Ident(lexeme.clone()),
                        Rule::INTEGER_CONST => TokenType::IntegerConst(parse_digits(&lexeme)),
                        _ => continue,
                    };
                    tokens.push(Token {
                        token_type,
                        lexeme,
                        line,
                    });
                }
            }
            Rule::ERROR_TOKEN => {
                let message = ErrorMessage {
                    line: pair.as_span().start_pos().line_col().0,
                    message: format!("Unrecognized token '{}'", pair.as_span().as_str()),
                };
                errors.push(message);
            }
            _ => {}
        }
    }
    if errors.is_empty() {
        return Ok(tokens);
    }
    Err(LexerError { messages: errors })
}

fn parse_digits(s: &str) -> usize {
    if s.starts_with("0x") || s.starts_with("0X") {
        usize::from_str_radix(&s[2..], 16).unwrap()
    } else if s.starts_with('0') && s.len() > 1 {
        usize::from_str_radix(&s[1..], 8).unwrap()
    } else {
        s.parse::<usize>().unwrap()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn tokenize_and_stringify(input: &str) -> String {
        match tokenize(input) {
            Ok(tokens) => {
                tokens
                    .iter()
                    .map(|t| t.to_string())
                    .collect::<Vec<String>>()
                    .join("\n")
                    + "\n"
            }
            Err(e) => format!("{}", e),
        }
    }

    #[test]
    fn test_case_1() {
        assert_eq!(
            tokenize_and_stringify(
                "int main()
{
   // line comment
   /*
     block comment
   */
   int i = 0x1;
}"
            ),
            "INT int at Line 1.
IDENT main at Line 1.
L_PAREN ( at Line 1.
R_PAREN ) at Line 1.
L_BRACE { at Line 2.
INT int at Line 7.
IDENT i at Line 7.
ASSIGN = at Line 7.
INTEGER_CONST 1 at Line 7.
SEMICOLON ; at Line 7.
R_BRACE } at Line 8.\n"
                .to_string()
        );
    }

    #[test]
    fn test_case_2() {
        assert_eq!(
            tokenize_and_stringify(
                "int main(){
  int i = 1;
  int j = ~i;
}"
            ),
            "Error type A at Line 3: Unrecognized token '~'.\n".to_string()
        );
    }

    #[test]
    fn test_case_3() {
        assert_eq!(
            tokenize_and_stringify(
                "int func(int arg) {
    int l;
    l = - - - arg;
    return l;
}

int main() {
    int x, y;
    x = 02;
    y = 0x1;
    x = x - 1 + y;
    if (+-!!!x) {
        x = - - -2;
    }
    else {
        x = 1 + + y;
    }
    func(x);
    return 0;
}"
            ),
            "INT int at Line 1.
IDENT func at Line 1.
L_PAREN ( at Line 1.
INT int at Line 1.
IDENT arg at Line 1.
R_PAREN ) at Line 1.
L_BRACE { at Line 1.
INT int at Line 2.
IDENT l at Line 2.
SEMICOLON ; at Line 2.
IDENT l at Line 3.
ASSIGN = at Line 3.
MINUS - at Line 3.
MINUS - at Line 3.
MINUS - at Line 3.
IDENT arg at Line 3.
SEMICOLON ; at Line 3.
RETURN return at Line 4.
IDENT l at Line 4.
SEMICOLON ; at Line 4.
R_BRACE } at Line 5.
INT int at Line 7.
IDENT main at Line 7.
L_PAREN ( at Line 7.
R_PAREN ) at Line 7.
L_BRACE { at Line 7.
INT int at Line 8.
IDENT x at Line 8.
COMMA , at Line 8.
IDENT y at Line 8.
SEMICOLON ; at Line 8.
IDENT x at Line 9.
ASSIGN = at Line 9.
INTEGER_CONST 2 at Line 9.
SEMICOLON ; at Line 9.
IDENT y at Line 10.
ASSIGN = at Line 10.
INTEGER_CONST 1 at Line 10.
SEMICOLON ; at Line 10.
IDENT x at Line 11.
ASSIGN = at Line 11.
IDENT x at Line 11.
MINUS - at Line 11.
INTEGER_CONST 1 at Line 11.
PLUS + at Line 11.
IDENT y at Line 11.
SEMICOLON ; at Line 11.
IF if at Line 12.
L_PAREN ( at Line 12.
PLUS + at Line 12.
MINUS - at Line 12.
NOT ! at Line 12.
NOT ! at Line 12.
NOT ! at Line 12.
IDENT x at Line 12.
R_PAREN ) at Line 12.
L_BRACE { at Line 12.
IDENT x at Line 13.
ASSIGN = at Line 13.
MINUS - at Line 13.
MINUS - at Line 13.
MINUS - at Line 13.
INTEGER_CONST 2 at Line 13.
SEMICOLON ; at Line 13.
R_BRACE } at Line 14.
ELSE else at Line 15.
L_BRACE { at Line 15.
IDENT x at Line 16.
ASSIGN = at Line 16.
INTEGER_CONST 1 at Line 16.
PLUS + at Line 16.
PLUS + at Line 16.
IDENT y at Line 16.
SEMICOLON ; at Line 16.
R_BRACE } at Line 17.
IDENT func at Line 18.
L_PAREN ( at Line 18.
IDENT x at Line 18.
R_PAREN ) at Line 18.
SEMICOLON ; at Line 18.
RETURN return at Line 19.
INTEGER_CONST 0 at Line 19.
SEMICOLON ; at Line 19.
R_BRACE } at Line 20.\n"
                .to_string()
        );
    }

    #[test]
    fn test_case_4() {
        assert_eq!(
            tokenize_and_stringify(
                "int array()
{
    int arr[10] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};

    int a1 = 0, a2 = 3, a3 = 5, a4 = 7, a5 = 9, a6 = 1, a7 = 2, a8 = 4,
        a9 = 6;

    return arr[a1] + arr[a2] + arr[a3] + arr[a4] + arr[a7] + arr[a8];
}

int main()
{
    int q = 1, r = 2, s = 04, t = 0x7, u = 0xA, v = 0xb, w = 0xcD, x = 077;

    int sum1 = q + r + s + t + u + v + w + x;

    int sum2 = array();

    int sum3 = sum1 + sum2;

    return 0;
}"
            ),
            "INT int at Line 1.
IDENT array at Line 1.
L_PAREN ( at Line 1.
R_PAREN ) at Line 1.
L_BRACE { at Line 2.
INT int at Line 3.
IDENT arr at Line 3.
L_BRACKT [ at Line 3.
INTEGER_CONST 10 at Line 3.
R_BRACKT ] at Line 3.
ASSIGN = at Line 3.
L_BRACE { at Line 3.
INTEGER_CONST 0 at Line 3.
COMMA , at Line 3.
INTEGER_CONST 1 at Line 3.
COMMA , at Line 3.
INTEGER_CONST 2 at Line 3.
COMMA , at Line 3.
INTEGER_CONST 3 at Line 3.
COMMA , at Line 3.
INTEGER_CONST 4 at Line 3.
COMMA , at Line 3.
INTEGER_CONST 5 at Line 3.
COMMA , at Line 3.
INTEGER_CONST 6 at Line 3.
COMMA , at Line 3.
INTEGER_CONST 7 at Line 3.
COMMA , at Line 3.
INTEGER_CONST 8 at Line 3.
COMMA , at Line 3.
INTEGER_CONST 9 at Line 3.
R_BRACE } at Line 3.
SEMICOLON ; at Line 3.
INT int at Line 5.
IDENT a1 at Line 5.
ASSIGN = at Line 5.
INTEGER_CONST 0 at Line 5.
COMMA , at Line 5.
IDENT a2 at Line 5.
ASSIGN = at Line 5.
INTEGER_CONST 3 at Line 5.
COMMA , at Line 5.
IDENT a3 at Line 5.
ASSIGN = at Line 5.
INTEGER_CONST 5 at Line 5.
COMMA , at Line 5.
IDENT a4 at Line 5.
ASSIGN = at Line 5.
INTEGER_CONST 7 at Line 5.
COMMA , at Line 5.
IDENT a5 at Line 5.
ASSIGN = at Line 5.
INTEGER_CONST 9 at Line 5.
COMMA , at Line 5.
IDENT a6 at Line 5.
ASSIGN = at Line 5.
INTEGER_CONST 1 at Line 5.
COMMA , at Line 5.
IDENT a7 at Line 5.
ASSIGN = at Line 5.
INTEGER_CONST 2 at Line 5.
COMMA , at Line 5.
IDENT a8 at Line 5.
ASSIGN = at Line 5.
INTEGER_CONST 4 at Line 5.
COMMA , at Line 5.
IDENT a9 at Line 6.
ASSIGN = at Line 6.
INTEGER_CONST 6 at Line 6.
SEMICOLON ; at Line 6.
RETURN return at Line 8.
IDENT arr at Line 8.
L_BRACKT [ at Line 8.
IDENT a1 at Line 8.
R_BRACKT ] at Line 8.
PLUS + at Line 8.
IDENT arr at Line 8.
L_BRACKT [ at Line 8.
IDENT a2 at Line 8.
R_BRACKT ] at Line 8.
PLUS + at Line 8.
IDENT arr at Line 8.
L_BRACKT [ at Line 8.
IDENT a3 at Line 8.
R_BRACKT ] at Line 8.
PLUS + at Line 8.
IDENT arr at Line 8.
L_BRACKT [ at Line 8.
IDENT a4 at Line 8.
R_BRACKT ] at Line 8.
PLUS + at Line 8.
IDENT arr at Line 8.
L_BRACKT [ at Line 8.
IDENT a7 at Line 8.
R_BRACKT ] at Line 8.
PLUS + at Line 8.
IDENT arr at Line 8.
L_BRACKT [ at Line 8.
IDENT a8 at Line 8.
R_BRACKT ] at Line 8.
SEMICOLON ; at Line 8.
R_BRACE } at Line 9.
INT int at Line 11.
IDENT main at Line 11.
L_PAREN ( at Line 11.
R_PAREN ) at Line 11.
L_BRACE { at Line 12.
INT int at Line 13.
IDENT q at Line 13.
ASSIGN = at Line 13.
INTEGER_CONST 1 at Line 13.
COMMA , at Line 13.
IDENT r at Line 13.
ASSIGN = at Line 13.
INTEGER_CONST 2 at Line 13.
COMMA , at Line 13.
IDENT s at Line 13.
ASSIGN = at Line 13.
INTEGER_CONST 4 at Line 13.
COMMA , at Line 13.
IDENT t at Line 13.
ASSIGN = at Line 13.
INTEGER_CONST 7 at Line 13.
COMMA , at Line 13.
IDENT u at Line 13.
ASSIGN = at Line 13.
INTEGER_CONST 10 at Line 13.
COMMA , at Line 13.
IDENT v at Line 13.
ASSIGN = at Line 13.
INTEGER_CONST 11 at Line 13.
COMMA , at Line 13.
IDENT w at Line 13.
ASSIGN = at Line 13.
INTEGER_CONST 205 at Line 13.
COMMA , at Line 13.
IDENT x at Line 13.
ASSIGN = at Line 13.
INTEGER_CONST 63 at Line 13.
SEMICOLON ; at Line 13.
INT int at Line 15.
IDENT sum1 at Line 15.
ASSIGN = at Line 15.
IDENT q at Line 15.
PLUS + at Line 15.
IDENT r at Line 15.
PLUS + at Line 15.
IDENT s at Line 15.
PLUS + at Line 15.
IDENT t at Line 15.
PLUS + at Line 15.
IDENT u at Line 15.
PLUS + at Line 15.
IDENT v at Line 15.
PLUS + at Line 15.
IDENT w at Line 15.
PLUS + at Line 15.
IDENT x at Line 15.
SEMICOLON ; at Line 15.
INT int at Line 17.
IDENT sum2 at Line 17.
ASSIGN = at Line 17.
IDENT array at Line 17.
L_PAREN ( at Line 17.
R_PAREN ) at Line 17.
SEMICOLON ; at Line 17.
INT int at Line 19.
IDENT sum3 at Line 19.
ASSIGN = at Line 19.
IDENT sum1 at Line 19.
PLUS + at Line 19.
IDENT sum2 at Line 19.
SEMICOLON ; at Line 19.
RETURN return at Line 21.
INTEGER_CONST 0 at Line 21.
SEMICOLON ; at Line 21.
R_BRACE } at Line 22.\n"
                .to_string()
        );
    }
}
